import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import re
from datetime import datetime

# Regular expression to extract accuracy
pattern = re.compile(r"Epoch \d+/\d+, loss: [\d.]+, accuracy: ([\d.]+)")


# Function to extract timestamps and compute training throughput
def extract_throughput(log_file_path, samples_per_epoch=25000):
    times = []
    epoch_durations = []
    epoch_throughputs = []

    with open(log_file_path, 'r') as file:
        start_time = None
        for line in file:
            # Extract timestamp
            time_match = re.search(r'\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}', line)
            if time_match:
                timestamp = datetime.strptime(time_match.group(), '%Y-%m-%d %H:%M:%S')
                times.append(timestamp)

                if start_time is None:
                    start_time = timestamp
                else:
                    duration = (timestamp - start_time).total_seconds()
                    epoch_durations.append(duration)
                    throughput = samples_per_epoch / duration
                    epoch_throughputs.append(throughput)
                    start_time = timestamp  # Reset for next epoch

    return epoch_throughputs


# Log file paths for all methods
log_files = {
    'NetSenseML-40G': 'vgg16_adp_40g.log',
    'AllReduce-40G': 'vgg16_allreduce_40g.log',
    'TopK-0.01-40G': 'vgg16_topk_40g.log',
    'NetSenseML-1G': 'vgg16_adpt_1g.log',
    'AllReduce-1G': 'vgg16_allreduce_1g.log',
}

# Extract training throughput for each method
throughputs = {}
for label, log_file in log_files.items():
    throughputs[label] = extract_throughput(log_file)


# Calculate average throughput and standard deviation for each method
avg_throughputs = [np.mean(throughputs[label]) for label in log_files.keys()]
std_devs = [np.std(throughputs[label]) for label in log_files.keys()]
print(avg_throughputs)
# Create the bar chart
labels = list(log_files.keys())
x_pos = np.arange(len(labels))

# Create figure and axis
fig, ax = plt.subplots(figsize=(10, 8))

# Define custom colors for each method
colors = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd']

# Plot bar chart with error bars
bars = ax.bar(x_pos, avg_throughputs, yerr=std_devs, align='center', alpha=0.85,
              capsize=18, color=colors, edgecolor='black', linewidth=1.2, error_kw={'elinewidth': 2, 'capthick': 2})

# Set tick positions and labels
ax.set_xticks(x_pos)
ax.set_xticklabels(labels, rotation=45, ha="right", fontsize=18)

# Title and labels
# ax.set_title('Training Throughput Comparison (samples/sec)', fontsize=18)
ax.set_xlabel('Method', fontsize=18)
ax.set_ylabel('Throughput (samples/sec)', fontsize=18)
ax.tick_params(axis='x', labelsize=18)
ax.tick_params(axis='y', labelsize=18)

# Grid for readability
ax.grid(True, which='major', linestyle='--', linewidth=0.7, axis='y', color='gray')

# Tight layout to prevent labels from being cut off
plt.tight_layout()

# Save the figure
plt.savefig('training_throughput_comparison.pdf', format='pdf', bbox_inches='tight')

# Show plot
plt.show()